# Copyright (c) 2008 Alvin Wang http://cloudwizard.com/
# License New BSD
# See README for additional information

"""Friendly Python SSH2 interface."""
#

import os
import tempfile
import paramiko
from paramiko import AutoAddPolicy
from boto.mashups import interactive
import tempfile
from cloudwizard.template import CWTemplate

class Connection(object):
    """Connects and logs into the specified hostname. 
    Arguments that are not given are guessed from the environment.""" 
    data={}
    def __init__(self,
                 host,
                 username = 'root',
                 private_key = None,
                 password = None,
                 port = 22,
                 data=None
                 ):
        self._sftp_live = False
        self._sftp = None
        if not username:
            username = os.environ['LOGNAME']

        # Log to a temporary file.
        templog = tempfile.mkstemp('.txt', 'ssh-')[1]
        paramiko.util.log_to_file(templog)

        # Begin the SSH transport.
        

        # need to add autoadd ****************************************
        self._transport = paramiko.SSHClient()
        self._transport.load_system_host_keys()
        self._transport.set_missing_host_key_policy(AutoAddPolicy())
        
        #self._tranport_live = True
        # Authenticate the transport.
        if password:
            # Using Password.
            self._transport.connect(host, port,username = username, password = password)
        else:
            # Use Private Key.
            if not private_key:
                # Try to use default key.
                if os.path.exists(os.path.expanduser('~/.ssh/id_rsa')):
                    private_key = '~/.ssh/id_rsa'
                elif os.path.exists(os.path.expanduser('~/.ssh/id_dsa')):
                    private_key = '~/.ssh/id_dsa'
                else:
                    raise TypeError, "You have not specified a password or key."

            private_key_file = os.path.expanduser(private_key)
            rsa_key = paramiko.RSAKey.from_private_key_file(private_key_file)
            self._transport.connect(host, port,username = username, pkey = rsa_key)
    
    def _sftp_connect(self):
        """Establish the SFTP connection."""
        if not self._sftp_live:
            self._sftp = paramiko.SFTPClient.from_transport(self._transport)
            self._sftp_live = True

    def get(self, remotepath, localpath = None):
        """Copies a file between the remote host and the local host."""
        if not localpath:
            localpath = os.path.split(remotepath)[1]
        self._sftp_connect()
        self._sftp.get(remotepath, localpath)

    def put(self, localpath=None, remotepath=None, url=None):
        """Copies a file between the local host and the remote host."""
        # if no localpath, get from URL
        if not remotepath:
            remotepath = os.path.split(localpath)[1]
        self._sftp_connect()
        self._sftp.put(localpath, remotepath)

    def execute(self, command):
        """Execute the given commands on a remote machine."""
        #channel = self._transport.open_session
        #paramiko.SFTPClient.from_transport(self._transport)
        stdin, stdout, stderr = self._transport.exec_command(command)
        
        result = stdout.readlines()
        return result

    def tput(self, localpath=None, remotepath=None, url=None, data=None):
        """Copies a file between the local host and the remote host."""
        # if no localpath, get from URL
        tpl = CWTemplate()
        tempfile = tpl.process( template_file=localpath, data=data)
        if not remotepath:
            remotepath = os.path.split(localpath)[1]
        self._sftp_connect()
        self._sftp.put( newfile, remotepath)

        os.remove(tempfile)

    def shell(self):
        """Execute the given commands on a remote machine."""
        print "<ctrl>Z when you are done"
        channel = self._transport.invoke_shell()
        interactive.interactive_shell(channel)
        # SSHClient.invoke_shell(self, term='vt100', width=80, height=24) 
                
    def close(self):
        """Closes the connection and cleans up."""
        # Close SFTP Connection.
        if self._sftp_live:
            self._sftp.close()
            self._sftp_live = False
        # Close the SSH Transport.
        self._transport.close()
        self._tranport_live = False

    def __del__(self):
        """Attempt to clean up if not explicitly closed."""
        self.close()
        
    def help(self):
        help(self)
        
    def get_distro(self):
        """ get file /etc/lsb-release
        """
        pass
